#!/user/bin/python
######################
## rebin histogram ###
## Lailin, 2012 ######
######################

import sys
ROOTSYS = '/afs/atlas.umich.edu/opt/root/lib'
sys.path.append(ROOTSYS)

#####################
## Import Module  ###
#####################
import array
import os
import glob
from math import sqrt,fabs,sin,log
from ROOT import TFile,TTree,TChain,TBranch,TH1,TH1F,TH2F,TList
from ROOT import TLorentzVector,TGraphAsymmErrors,TMath
from ROOT import THStack,TCanvas,TLegend,TColor,TPaveText,TPad
from ROOT import gStyle,gDirectory
from ROOT import Double

def get_hist_2D_significance(histsig='', histbkg='', cutdirectionx=1, cutdirectiony=1, profile=1, **kw):
    '''
    get siginificane using 2D scan 
    cutdirection: 1: Integral(i, nbins); 2: Integral(i, nbins+1)
                 -1: Integral(1, i);    -2: Integral(0, i)
    '''

    ## book histogram
    nbinsx=histsig.GetNbinsX()
    nbinsy=histsig.GetNbinsY()
    histo_xmin=histsig.GetXaxis().GetXmin()
    histo_xmax=histsig.GetXaxis().GetXmax()
    histo_ymin=histsig.GetYaxis().GetXmin()
    histo_ymax=histsig.GetYaxis().GetXmax()
    hist_signf=TH2F('significance', 'significance', nbinsx, histo_xmin, histo_xmax, nbinsy, histo_ymin, histo_ymax)
    hist_signf.Sumw2()
    for i in range(1,nbinsx+1):
      for j in range(1,nbinsy+1):
        nsig, nbkg=0, 0
        ## cut is like: xx > xx
        if cutdirectionx==1 and cutdirectiony==1: 
            nsig=histsig.Integral(i, nbinsx, j, nbinsy)  # no overflow
            nbkg=histbkg.Integral(i, nbinsx, j, nbinsy)  # no overflow
        elif cutdirectionx==1 and cutdirectiony==2: 
            nsig=histsig.Integral(i, nbinsx, j, nbinsy+1)  
            nbkg=histbkg.Integral(i, nbinsx, j, nbinsy+1)  
        elif cutdirectionx==2 and cutdirectiony==1: 
            nsig=histsig.Integral(i, nbinsx+1, j, nbinsy)  
            nbkg=histbkg.Integral(i, nbinsx+1, j, nbinsy)  
        elif cutdirectionx==2 and cutdirectiony==2: 
            nsig=histsig.Integral(i, nbinsx+1, j, nbinsy+1)
            nbkg=histbkg.Integral(i, nbinsx+1, j, nbinsy+1)
        elif cutdirectionx==-2 and cutdirectiony==-2: 
            nsig=histsig.Integral(0, i, 0, j)
            nbkg=histbkg.Integral(0, i, 0, j)
        elif cutdirectionx==-2 and cutdirectiony==-1: 
            nsig=histsig.Integral(0, i, 1, j)
            nbkg=histbkg.Integral(0, i, 1, j)
        elif cutdirectionx==-1 and cutdirectiony==-2: 
            nsig=histsig.Integral(1, i, 0, j)
            nbkg=histbkg.Integral(1, i, 0, j)
        elif cutdirectionx==-1 and cutdirectiony==-1: 
            nsig=histsig.Integral(1, i, 1, j) # no underflow
            nbkg=histbkg.Integral(1, i, 1, j) # no underflow
        elif cutdirectionx==1 and cutdirectiony==-1: 
            nsig=histsig.Integral(i, nbinsx, 1, j)
            nbkg=histbkg.Integral(i, nbinsx, 1, j)
        elif cutdirectionx==2 and cutdirectiony==-1: 
            nsig=histsig.Integral(i, nbinsx+1, 1, j)
            nbkg=histbkg.Integral(i, nbinsx+1, 1, j)
        elif cutdirectionx==1 and cutdirectiony==-2: 
            nsig=histsig.Integral(i, nbinsx, 0, j)
            nbkg=histbkg.Integral(i, nbinsx, 0, j)
        elif cutdirectionx==2 and cutdirectiony==2: 
            nsig=histsig.Integral(i, nbinsx+1, 0, j)
            nbkg=histbkg.Integral(i, nbinsx+1, 0, j)
        elif cutdirectionx==-1 and cutdirectiony==1: 
            nsig=histsig.Integral(1, i, j, nbinsy)
            nbkg=histbkg.Integral(1, i, j, nbinsy)
        elif cutdirectionx==-2 and cutdirectiony==1: 
            nsig=histsig.Integral(0, i, j, nbinsy)
            nbkg=histbkg.Integral(0, i, j, nbinsy)
        elif cutdirectionx==-1 and cutdirectiony==2: 
            nsig=histsig.Integral(1, i, j, nbinsy+1)
            nbkg=histbkg.Integral(1, i, j, nbinsy+1)
        elif cutdirectionx==-2 and cutdirectiony==2: 
            nsig=histsig.Integral(0, i, j, nbinsy+1)
            nbkg=histbkg.Integral(0, i, j, nbinsy+1)
        #print "sig/bkg: ", nsig, nbkg
        signf=get_signf(nsig, nbkg, profile)
        hist_signf.SetBinContent(i, j, signf)

    return hist_signf

def get_signf(nsig=1, nbkg=1, profile=1,**kw):
    signf=0
    if profile==0:
        if nbkg!=0: signf=nsig/sqrt(nbkg) 
        else: signf=-1
    if profile==1:
        if nbkg!=0: signf=sqrt(2*((nsig+nbkg)*log(1+nsig/nbkg) - nsig))
        else: signf=-1
    return signf
